#!/usr/bin/env python3
"""
Prepare script: estimate .npz sizes & extraction time, pre-create _memmap, and GPU self-check.

Usage examples:
  python scripts/prepare_memmap_and_gpu_check.py \
    --base ./data/SINPA --device cuda

  python scripts/prepare_memmap_and_gpu_check.py \
    --base D:/datasets/SINPA --memmap-dir D:/datasets/SINPA/_memmap --device cuda

"""

import os
import sys
import time
import argparse
import zipfile
from typing import List, Tuple

import numpy as np


def estimate_npz(npz_path: str, mbps_list: List[int]) -> int:
    """Estimate uncompressed sizes of x.npy/y.npy in npz and time by throughput.
    Returns total bytes of x+y (uncompressed)."""
    if not os.path.exists(npz_path):
        print(f"[ESTIMATE] Missing npz: {npz_path}")
        return 0
    try:
        with zipfile.ZipFile(npz_path, 'r') as zf:
            names = zf.namelist()
            x_name = next((n for n in names if n.endswith('x.npy')), None)
            y_name = next((n for n in names if n.endswith('y.npy')), None)
            x_size = zf.getinfo(x_name).file_size if x_name else 0
            y_size = zf.getinfo(y_name).file_size if y_name else 0
            total = x_size + y_size
            print(f"[ESTIMATE] {npz_path}")
            print(f"           x.npy: {x_size/1024/1024:.2f} MB, y.npy: {y_size/1024/1024:.2f} MB, total: {total/1024/1024:.2f} MB")
            for mbps in mbps_list:
                if mbps <= 0:
                    continue
                sec = total / (mbps * 1024 * 1024)
                print(f"           ~{mbps} MB/s => ~{sec:.1f} s")
            return total
    except Exception as e:
        print(f"[ESTIMATE] Failed to read {npz_path}: {e}")
        return 0


def ensure_memmap(npz_path: str, memmap_dir: str, split: str) -> Tuple[str, str, tuple, tuple]:
    """Invoke project helper to extract x.npy/y.npy into memmap dir and validate shapes."""
    from src.utils.helper import _ensure_memmap_paths
    os.makedirs(memmap_dir, exist_ok=True)
    x_path, y_path = _ensure_memmap_paths(npz_path, memmap_dir, split)
    print(f"[MEMMAP] {split}: x -> {x_path}, y -> {y_path}")
    x = np.load(x_path, mmap_mode='r')
    y = np.load(y_path, mmap_mode='r')
    print(f"[MEMMAP] {split} shapes: x={getattr(x,'shape',None)}, y={getattr(y,'shape',None)}")
    return x_path, y_path, getattr(x, 'shape', None), getattr(y, 'shape', None)


def gpu_self_check(device_str: str = 'cuda') -> bool:
    """Check torch+CUDA availability and run a small matmul benchmark."""
    try:
        import torch
    except Exception as e:
        print(f"[GPU] PyTorch import failed: {e}")
        return False

    use_cuda = device_str.startswith('cuda') and torch.cuda.is_available()
    dev = torch.device('cuda') if use_cuda else torch.device('cpu')
    print(f"[GPU] torch: {torch.__version__}, cuda_available: {torch.cuda.is_available()}")
    if dev.type == 'cuda':
        try:
            print(f"[GPU] device: {torch.cuda.get_device_name(dev)}")
        except Exception:
            pass
    # simple compute test
    try:
        t0 = time.time()
        a = torch.randn((2048, 2048), device=dev)
        b = torch.randn((2048, 2048), device=dev)
        c = a @ b
        if dev.type == 'cuda':
            torch.cuda.synchronize()
        t1 = time.time()
        print(f"[GPU] matmul 2048x2048 time: {(t1 - t0):.3f} s on {dev.type}")
        return True
    except Exception as e:
        print(f"[GPU] compute test failed: {e}")
        return False


def main():
    parser = argparse.ArgumentParser(description="Pre-extract _memmap from npz and check GPU")
    parser.add_argument('--base', type=str, default='./data/SINPA', help='Dataset base directory (contains train/val/test npz)')
    parser.add_argument('--memmap-dir', type=str, default=None, help='Output memmap directory (default: <base>/_memmap)')
    parser.add_argument('--splits', type=str, nargs='*', default=['train', 'val', 'test'], help='Splits to process')
    parser.add_argument('--estimate-mbps', type=int, nargs='*', default=[50, 100, 200, 300], help='Throughput list for time estimates')
    parser.add_argument('--device', type=str, default='cuda', choices=['cuda', 'cpu'], help='Device for GPU self-check')
    parser.add_argument('--skip-extract', action='store_true', help='Only estimate sizes; do not extract to _memmap')

    args = parser.parse_args()

    base = os.path.abspath(args.base)
    memdir = args.memmap_dir or os.path.join(base, '_memmap')

    print(f"[INFO] base: {base}")
    print(f"[INFO] memmap_dir: {memdir}")
    print(f"[INFO] splits: {args.splits}")

    # 1) GPU self-check
    gpu_ok = gpu_self_check(args.device)
    if not gpu_ok:
        print("[GPU] Warning: GPU self-check did not complete successfully. Training may fall back to CPU.")

    # 2) Estimate sizes & optional extract
    total_bytes_all = 0
    for split in args.splits:
        npz_path = os.path.join(base, f"{split}.npz")
        total_bytes_all += estimate_npz(npz_path, args.estimate_mbps)

    if args.skip_extract:
        print("[MEMMAP] Skipped extraction per --skip-extract")
    else:
        for split in args.splits:
            npz_path = os.path.join(base, f"{split}.npz")
            if not os.path.exists(npz_path):
                print(f"[MEMMAP] Skip {split}: missing {npz_path}")
                continue
            try:
                ensure_memmap(npz_path, memdir, split)
            except Exception as e:
                print(f"[MEMMAP] {split} extraction failed: {e}")

    print("[DONE] Preparation completed.")
    if total_bytes_all > 0:
        print(f"[DONE] Estimated total uncompressed x+y: {total_bytes_all/1024/1024:.2f} MB across splits.")
        print("[NEXT] You can now run training with fast memmap path (if extracted).")


if __name__ == '__main__':
    main()